package com.sankuai.inf.leaf.segment.dao.impl;

import com.sankuai.inf.leaf.segment.dao.IDAllocDao;
import com.sankuai.inf.leaf.segment.model.LeafAlloc;
import lombok.extern.slf4j.Slf4j;
import xyz.jcat.common.exception.SystemException;
import xyz.jcat.common.util.CollectionUtils;
import xyz.jcat.common.util.DateUtils;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.text.MessageFormat;
import java.util.ArrayList;
import java.util.Date;
import java.util.List;
import java.util.stream.Collectors;

@Slf4j
public class IDAllocDaoImpl implements IDAllocDao {

    private DataSource dataSource;

    private static final String SQL_GET_ALL_LEAF_ALLOCS = "SELECT biz_tag, max_id, step, update_time FROM leaf_alloc";
    private static final String SQL_UPDATE_MAX_ID = "UPDATE leaf_alloc SET max_id = max_id + step, update_time = ''{0}'' WHERE biz_tag = ''{1}''";
    private static final String SQL_UPDATE = "UPDATE leaf_alloc SET max_id = max_id + {0}, update_time = ''{1}'' WHERE biz_tag = ''{2}''";
    private static final String SQL_GET_LEAF_ALLOC = "SELECT biz_tag, max_id, step, update_time FROM leaf_alloc WHERE biz_tag = ''{0}''";

    @Override
    public List<LeafAlloc> getAllLeafAllocs() {
        List<LeafAlloc> list = new ArrayList<>();
        try (Connection conn = dataSource.getConnection();
             PreparedStatement ps = conn.prepareStatement(SQL_GET_ALL_LEAF_ALLOCS);
             ResultSet rst = ps.executeQuery();
        ){
            while (rst.next()) {
                list.add(getLeafAlloc(rst));
            }
        }catch (Exception e) {
            throw new SystemException(e);
        }
        return list;
    }

    @Override
    public LeafAlloc updateMaxIdAndGetLeafAlloc(String tag) {
        updateMaxId(tag);
        return getLeafAlloc(tag);
    }

    private void updateMaxId(String tag) {
        String sql = MessageFormat.format(SQL_UPDATE_MAX_ID, DateUtils.format(new Date()), tag);
        try (Connection conn = dataSource.getConnection();
             PreparedStatement ps = conn.prepareStatement(sql);
        ){
            ps.executeUpdate();
        }catch (Exception e) {
            throw new SystemException(e);
        }
    }

    private void update(LeafAlloc leafAlloc) {
        String sql = MessageFormat.format(SQL_UPDATE, leafAlloc.getStep(), DateUtils.format(new Date()), leafAlloc.getKey());
        try (Connection conn = dataSource.getConnection();
             PreparedStatement ps = conn.prepareStatement(sql);
        ){
            ps.executeUpdate();
        }catch (Exception e) {
            throw new SystemException(e);
        }
    }

    private LeafAlloc getLeafAlloc(String tag) {
        String sql = MessageFormat.format(SQL_GET_LEAF_ALLOC, tag);
        try (Connection conn = dataSource.getConnection();
             PreparedStatement ps = conn.prepareStatement(sql);
             ResultSet rst = ps.executeQuery();
        ){
            if(rst.next()) {
                return getLeafAlloc(rst);
            }
        }catch (Exception e) {
            throw new SystemException(e);
        }
        return null;
    }

    private LeafAlloc getLeafAlloc(ResultSet rst) throws SQLException {
        LeafAlloc leafAlloc = new LeafAlloc();
        leafAlloc.setKey(rst.getString("biz_tag"));
        leafAlloc.setMaxId(rst.getLong("max_id"));
        leafAlloc.setStep(rst.getInt("step"));
        leafAlloc.setUpdateTime(rst.getString("update_time"));
        return leafAlloc;
    }

    @Override
    public LeafAlloc updateMaxIdByCustomStepAndGetLeafAlloc(LeafAlloc leafAlloc) {
        update(leafAlloc);
        return getLeafAlloc(leafAlloc.getKey());
    }

    @Override
    public List<String> getAllTags() {
        List<LeafAlloc> list = getAllLeafAllocs();
        if(CollectionUtils.isNotEmpty(list)) {
            return list.stream().map(LeafAlloc::getKey).collect(Collectors.toList());
        }
        return new ArrayList<>(0);
    }

    public IDAllocDaoImpl(DataSource dataSource) {
        this.dataSource = dataSource;
    }
}
